"""
PMMA Slab exposed to uniform surface heat flux

M. Bruns
April 2015
"""

from dolfin import *
from mshr import *
import numpy as np

parameters["allow_extrapolation"] = True

# ----------------------------------------------------------------------------
# Problem Parameters
# ----------------------------------------------------------------------------

# numerical parameters
N = 32
dt = 1e0
t_f = 120.

# scenario parameters
length = 0.080       # m 
length = length/2.
height = 8.5e-3      # m
q_ext_v = Constant((0.0, -52e3))    # W/m^2
h = 10.         # W/m^2-K
T_inf = 298.    # K

# material properties
k = 0.2         # W/m-K
rho = 1100.     # kg/m^3
c_p = 2200.     # J/kg-K
eps = 0.85
A = 8.5e12      # 1/s
E_a = 1.88e5    # J/mol
dh = 870e3      # J/kg
mu = 5e4        # N-s/m^2 (van Krevelen at 200 C)

# constants
R = 8.314       # J/K-mol
sigma = 5.67e-8 # W/m^2-K^4

# ----------------------------------------------------------------------------
# Supporting Functions and Classes
# ----------------------------------------------------------------------------

# Class for interfacing with the Newton solver
class NonlinearEquation(NonlinearProblem):
    def __init__(self, a, L):
        NonlinearProblem.__init__(self)
        self.L = L
        self.a = a
    def F(self, b, x):
        assemble(self.L, tensor=b)
    def J(self, A, x):
        assemble(self.a, tensor=A)

# boundary classes
class LeftBoundary(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and abs(x[0]) < DOLFIN_EPS

class BottomBoundary(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and abs(x[1]) < DOLFIN_EPS

class RightBoundary(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and abs(x[0] - length) < DOLFIN_EPS

class TopBoundary(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and abs(x[1] - height) < DOLFIN_EPS

# source term functions
def rxn_rate(T):
    # rate of reaction, 1/s
    return A*exp(-E_a/(R*T))

def q_rad(T):
    # radiation into sample
    return eps*sigma*(T_inf**4 - T**4)

def q_conv(T):
    # convection into sample
    return h*(T_inf - T)

# remesh subroutine
def remesh(mesh, N):
    
    mesh.init()
   
    # get indice pairs of boundary vertices
    edge_pairs = []
    
    for f in facets(mesh):
        if f.exterior():
            pair = []
            for v in vertices(f):
                pair.append(v.index())
            edge_pairs.append(sorted(pair))
   
    # sort edge pairs by first index
    def getKey(item):
        return item[0]
    
    edge_pairs = sorted(edge_pairs, key=getKey)
   
    # build ordered list of connected vertices
    ordered_indices = edge_pairs[0]
    edge_pairs.remove(edge_pairs[0])
    
    while len(edge_pairs) > 0:
        
        for pair in edge_pairs:
            
            if pair[0] == ordered_indices[-1]:
                ordered_indices.append(pair[1])
                edge_pairs.remove(pair)
                break
            
            if pair[1] == ordered_indices[-1]:
                ordered_indices.append(pair[0])
                edge_pairs.remove(pair)
                break
   
    # build list of points to construct bounding polygon
    mesh_points = []
   
    # ...add first point
    v = Vertex(mesh, ordered_indices[0])
    v_point = Point(v.x(0), v.x(1))
    mesh_points.append(v_point)
   
    # distance between included vertices
    d_max = height/15.

    for i in ordered_indices[1:-1]:
        v = Vertex(mesh, i)
        v_point = Point(v.x(0), v.x(1))
        d = v_point.distance(mesh_points[-1])
        if d > d_max:
            mesh_points.append(v_point)
        elif abs(v.x(0)) < DOLFIN_EPS and abs(v.x(1)) < DOLFIN_EPS:
            # include lower left vertex
            mesh_points.append(v_point)
        elif abs(v.x(0) - length) < DOLFIN_EPS and abs(v.x(1)) < DOLFIN_EPS:
            # include lower right vertex
            mesh_points.append(v_point)
   
    if Vertex(mesh, ordered_indices[1]).x(0) \
            <= Vertex(mesh, ordered_indices[0]).x(0) + DOLFIN_EPS:
        mesh_points.reverse()
        print "reversing"

    return generate_mesh(Polygon(mesh_points), N)

# update boundary markers subroutine
def update_boundaries(mesh):

    boundary = MeshFunction("size_t", mesh, mesh.topology().dim()-1)
   
    for f in facets(mesh):
        boundary[f.index()] = 999

    # boundary conditions
    bottom_boundary.mark(boundary, 1)
    left_boundary.mark(boundary, 3)
    right_boundary.mark(boundary, 1)
   
    for f in facets(mesh):
        if f.exterior() and boundary[f.index()] != 1 and boundary[f.index()] != 3:
            if f.normal(1) >= DOLFIN_EPS:
                boundary[f.index()] = 2
            else:
                boundary[f.index()] = 4

    return boundary


# ----------------------------------------------------------------------------
# Velocity Solver
# ----------------------------------------------------------------------------

def velocity(mesh, T):
    
    # define function spaces
    V = VectorFunctionSpace(mesh, "CG", 2)
    Q = FunctionSpace(mesh, "CG", 1)
    W = V * Q

    T = interpolate(T, Q)

    # boundary conditions
    
    # ...no slip (ns)
    bc1 = DirichletBC(W.sub(0), Constant((0.0,0.0)), bottom_boundary)
    bc2 = DirichletBC(W.sub(0).sub(0), Constant(0.0), left_boundary)
    bc3 = DirichletBC(W.sub(0), Constant((0.0,0.0)), right_boundary)
     
    bcs = [bc1, bc2, bc3]

    G = Constant((0.0, -9.81))

    # define variational problem
    (u, p) = TrialFunctions(W)
    (v, q) = TestFunctions(W)

    a = + inner(mu*grad(u), grad(v))*dx                 \
        - p*div(v)*dx                                   \
        - div(u)*q*dx                                   \
        - inner(p*normal, v)*dss(3)

    L = + rxn_rate(T)*q*dx                              \
        + inner(rho*G, v)*dx                            \
        - inner((mu/3.)*grad(rxn_rate(T)), v)*dx

    w = Function(W)

    A, b = assemble_system(a, L, bcs)
    
    # # lu decomposition
    # solver = LUSolver(A)
    # solver.solve(w.vector(), b)
    
    # minres
    solver = KrylovSolver("minres", "petsc_amg")
    solver.set_operator(A)
    solver.solve(w.vector(), b)

    (u, p) = w.split(True)

    return u

# ----------------------------------------------------------------------------
# Mesh and Boundary Conditions
# ----------------------------------------------------------------------------

# Create mesh and define function spaces
mesh = generate_mesh(Rectangle(dolfin.Point(0., 0.), 
                               dolfin.Point(length, height)), N)

print "Number of cells in mesh = ", mesh.num_cells()

V = FunctionSpace(mesh, "Lagrange", 1)

# boundary conditions
bottom_boundary = BottomBoundary()
left_boundary = LeftBoundary()
right_boundary = RightBoundary()
boundary = update_boundaries(mesh)

# ----------------------------------------------------------------------------
# Weak Statement of Problem
# ----------------------------------------------------------------------------

# Define trial and test functions
dT = TrialFunction(V)
v  = TestFunction(V)

# Define functions
T   = Function(V)  # current solution
T_0  = Function(V)  # solution from previous converged step
T_inf = Constant(T_inf)

# Weak statement of the equations
dss = ds(domain=mesh, subdomain_data=boundary)
normal = FacetNormal(mesh)
L = T*v*dx - T_0*v*dx + (k*dt/(rho*c_p))*inner(grad(T), grad(v))*dx         \
    + (dt*dh*rxn_rate(T)/c_p)*v*dx                                          \
    + (dt/(rho*c_p))*(dot(q_ext_v, normal) - q_conv(T) - q_rad(T))*v*dss(2) \
    + (dt/(rho*c_p))*(- q_conv(T) - q_rad(T))*v*dss(4)

# Compute directional derivative about u in the direction of du (Jacobian)
a = derivative(L, T, dT)

# ----------------------------------------------------------------------------
# Solver Setup
# ----------------------------------------------------------------------------

# Form compiler options
parameters["form_compiler"]["optimize"]     = True
parameters["form_compiler"]["cpp_optimize"] = True
parameters["form_compiler"]["representation"] = "quadrature"

# Create nonlinear problem and Newton solver
problem = NonlinearEquation(a, L)
solver = NewtonSolver()
solver.parameters["linear_solver"] = "lu"
solver.parameters["convergence_criterion"] = "incremental"
solver.parameters["relative_tolerance"] = 1e-6
solver.parameters["maximum_iterations"] = 10
solver.parameters["relaxation_parameter"] = 1.0
solver.parameters["absolute_tolerance"] = 1e-15
solver.parameters["report"]= False

# ----------------------------------------------------------------------------
# Solve
# ----------------------------------------------------------------------------

# Output file
file_temp = File("./output/paraview/temperatures.pvd", "compressed") 

# Create intial conditions and interpolate
T.interpolate(T_inf)
T_0.interpolate(T_inf)
U = VectorFunctionSpace(mesh, "Lagrange", 2)
u = Function(U)
disp = Function(U)

# set up MLR output file
file_mlr = './output/pmma_mlr.csv'
f_handle = open(file_mlr, 'w')
np.savetxt( f_handle, np.array([0., 0.]).reshape(1,2), fmt='%.5e',
            delimiter=',',
            header='time (s), MLR (kg/m^2-s)', comments='')
f_handle.close()

rad_ratio = 0.1

# Step in time
t = 0.0
while (t < t_f):
    
    t += dt
    # print status to standard out
    print "Time = ", t, ", s"
    
    T_0.vector()[:] = T.vector()
    solver.solve(problem, T.vector())

    # move mesh
    u = velocity(mesh, T)
    disp.assign(u*dt)
    mesh.move(disp)
    mesh.smooth(3)
    
    # update boundary
    boundary = update_boundaries(mesh)
    dss = ds(domain=mesh, subdomain_data=boundary)

    # save output
    file_temp << (T, t)

    # compute MLR and print to file
    MLR = rho*assemble(rxn_rate(T)*dx)/length
    f_handle = open(file_mlr, 'a')
    np.savetxt( f_handle, np.array([t, MLR]).reshape(1,2), fmt='%.5e',
                delimiter=',', comments='')
    f_handle.close()

    # remesh if necessary
    if (mesh.rmin()/mesh.rmax()) < rad_ratio:
       
        print "Remeshing..."
        
        # compute temperatures on new mesh
        mesh_new = remesh(mesh, N)
        V_new = FunctionSpace(mesh_new, "Lagrange", 1)
        T_new = interpolate(T, V_new)

        # update mesh and function space
        mesh = mesh_new
        boundary = update_boundaries(mesh)
        V = V_new
        
        # ----------------------------------------------------------------------------
        # Solver Setup
        # ----------------------------------------------------------------------------
   
        # Define trial and test functions
        dT = TrialFunction(V)
        v  = TestFunction(V)
        
        # Define functions
        T   = Function(V)  # current solution
        T_0  = Function(V)  # solution from previous converged step
        
        # Weak statement of the equations
        dss = ds(domain=mesh, subdomain_data=boundary)
        normal = FacetNormal(mesh)
        L = T*v*dx - T_0*v*dx + (k*dt/(rho*c_p))*inner(grad(T), grad(v))*dx         \
            + (dt*dh*rxn_rate(T)/c_p)*v*dx                                          \
            + (dt/(rho*c_p))*(dot(q_ext_v, normal) - q_conv(T) - q_rad(T))*v*dss(2) \
            + (dt/(rho*c_p))*(- q_conv(T) - q_rad(T))*v*dss(4)
        
        # Compute directional derivative about u in the direction of du (Jacobian)
        a = derivative(L, T, dT)
        
        # Create nonlinear problem and Newton solver
        problem = NonlinearEquation(a, L)
        solver = NewtonSolver()
        solver.parameters["linear_solver"] = "lu"
        solver.parameters["convergence_criterion"] = "incremental"
        solver.parameters["relative_tolerance"] = 1e-6
        solver.parameters["maximum_iterations"] = 10
        solver.parameters["relaxation_parameter"] = 1.0
        solver.parameters["absolute_tolerance"] = 1e-15
        solver.parameters["report"]= False

        # Reinitialize temperature
        T.vector()[:] = T_new.vector()

        # Reinitialize rate and velocity functions
        U = VectorFunctionSpace(mesh, "Lagrange", 2)
        u = Function(U)
        disp = Function(U)

